setwd("/Users/markjchin/Dropbox/iat_seda_ocr/iat_seda_ocr/")

# Log
sink(file="R/cr-estimate-mrp.txt") 

# Include necessary packages
rm(list = ls())
library(foreign)
library(tidyverse)
library(feather)
library(lme4)
library(lubridate)
library(RCurl)

### Create MRP scores, Teacher
  
  # Load Project Implicit Data 
  iat.df <- read.dta('data/iat_mrp_tch.dta')
  
  # Load synthetic county data (state by state)
  poststrat.df <- read.dta('data/teacher_synth_county.dta')

  # Estimate MRP Models
  iat.imp.mrp.1 <- lmer(implicitbias ~ 1 +
                          (1 | countyid) + 
                          (1 | fips:race),
                        data = iat.df, verbose=T, 
                        control=lmerControl(optimizer="bobyqa",
                                            optCtrl=list(maxfun=2e5)))
  
  saveRDS(iat.imp.mrp.1, 'R/mrp-imp-tch')
  
  iat.exp.mrp.1 <- lmer(explicitbias ~ 1 +
                          (1 | countyid) + 
                          (1 | fips:race),
                        data = iat.df, verbose=T, 
                        control=lmerControl(optimizer="bobyqa",
                                            optCtrl=list(maxfun=2e5)))
  
  saveRDS(iat.exp.mrp.1, 'R/mrp-tch-exp')

  # Generate Poststratification Predictions 
  
  iat.imp.mrp1.poststrat.pred <- predict(iat.imp.mrp.1, 
                                         newdata=poststrat.df, 
                                         type='response',
                                         allow.new.levels=T)
  
  
  iat.exp.mrp1.poststrat.pred <- predict(iat.exp.mrp.1, 
                                         newdata=poststrat.df, 
                                         type='response',
                                         allow.new.levels=T)
  
  iat.mrp1.poststrat.est <- base::cbind(poststrat.df, 
                                        iat.imp.mrp1.poststrat.pred = iat.imp.mrp1.poststrat.pred,
                                        iat.exp.mrp1.poststrat.pred = iat.exp.mrp1.poststrat.pred) %>%
    group_by(countyid, tottch) %>%
    summarize(mrp1.imp_bias_tch = sum(prop*iat.imp.mrp1.poststrat.pred),
              mrp1.exp_bias_tch = sum(prop*iat.exp.mrp1.poststrat.pred)) %>%
    ungroup() 
  
  write.dta(iat.mrp1.poststrat.est, "data/mrp_tch_scores.dta")

### Create MRP scores, Pooled
  
  rm(list = ls())
  
  # Load Project Implicit Data 
  iat.df <- read.dta('data/iat_mrp_all.dta')
  
  # Load synthetic county data (5 year panel 2010-2015)
  poststrat.df <- read.dta('data/synth_county.dta')

  # Esitmate models
  iat.imp.mrp.1 <- lmer(implicitbias ~ 1 +
                          (1 | countyid) + 
                          (1 | fips) +
                          (1 | division:race) + 
                          (1 | division:sex:age),
                        data = iat.df, verbose=T, 
                        control=lmerControl(optimizer="bobyqa",
                                            optCtrl=list(maxfun=2e5)))
  saveRDS(iat.imp.mrp.1, 'R/mrp-all-imp')
  
  iat.exp.mrp.1 <- lmer(explicitbias ~ 1 +
                          (1 | countyid) + 
                          (1 | fips) +
                          (1 | division:race) + 
                          (1 | division:sex:age),
                        data = iat.df, verbose=T, 
                        control=lmerControl(optimizer="bobyqa",
                                            optCtrl=list(maxfun=2e5)))
  saveRDS(iat.exp.mrp.1, 'R/mrp-all-exp')
  
  # Generate Poststratification Predictions 
  
  iat.imp.mrp1.poststrat.pred <- predict(iat.imp.mrp.1, 
                                         newdata=poststrat.df, 
                                         type='response',
                                         allow.new.levels=T)
  
  iat.exp.mrp1.poststrat.pred <- predict(iat.exp.mrp.1, 
                                         newdata=poststrat.df, 
                                         type='response',
                                         allow.new.levels=T)
  
  iat.mrp1.poststrat.est <- base::cbind(poststrat.df, 
                                        iat.imp.mrp1.poststrat.pred = iat.imp.mrp1.poststrat.pred,
                                        iat.exp.mrp1.poststrat.pred = iat.exp.mrp1.poststrat.pred) %>%
    group_by(countyid, totpop) %>%
    summarize(mrp1.imp_bias = sum(prop*iat.imp.mrp1.poststrat.pred),
              mrp1.exp_bias = sum(prop*iat.exp.mrp1.poststrat.pred)) %>%
    ungroup() 
  
  write.dta(iat.mrp1.poststrat.est, "data/mrp_all_scores.dta")

sink()
